{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using bos_token, but it is not set yet.\n",
      "Using sep_token, but it is not set yet.\n",
      "Using cls_token, but it is not set yet.\n",
      "Using mask_token, but it is not set yet.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ByT5Tokenizer(name_or_path='', vocab_size=256, model_max_length=1000000000000000019884624838656, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'eos_token': '</s>', 'unk_token': '<unk>', 'pad_token': '<pad>', 'additional_special_tokens': ['<extra_id_0>', '<extra_id_1>', '<extra_id_2>', '<extra_id_3>', '<extra_id_4>', '<extra_id_5>', '<extra_id_6>', '<extra_id_7>', '<extra_id_8>', '<extra_id_9>', '<extra_id_10>', '<extra_id_11>', '<extra_id_12>', '<extra_id_13>', '<extra_id_14>', '<extra_id_15>', '<extra_id_16>', '<extra_id_17>', '<extra_id_18>', '<extra_id_19>', '<extra_id_20>', '<extra_id_21>', '<extra_id_22>', '<extra_id_23>', '<extra_id_24>', '<extra_id_25>', '<extra_id_26>', '<extra_id_27>', '<extra_id_28>', '<extra_id_29>', '<extra_id_30>', '<extra_id_31>', '<extra_id_32>', '<extra_id_33>', '<extra_id_34>', '<extra_id_35>', '<extra_id_36>', '<extra_id_37>', '<extra_id_38>', '<extra_id_39>', '<extra_id_40>', '<extra_id_41>', '<extra_id_42>', '<extra_id_43>', '<extra_id_44>', '<extra_id_45>', '<extra_id_46>', '<extra_id_47>', '<extra_id_48>', '<extra_id_49>', '<extra_id_50>', '<extra_id_51>', '<extra_id_52>', '<extra_id_53>', '<extra_id_54>', '<extra_id_55>', '<extra_id_56>', '<extra_id_57>', '<extra_id_58>', '<extra_id_59>', '<extra_id_60>', '<extra_id_61>', '<extra_id_62>', '<extra_id_63>', '<extra_id_64>', '<extra_id_65>', '<extra_id_66>', '<extra_id_67>', '<extra_id_68>', '<extra_id_69>', '<extra_id_70>', '<extra_id_71>', '<extra_id_72>', '<extra_id_73>', '<extra_id_74>', '<extra_id_75>', '<extra_id_76>', '<extra_id_77>', '<extra_id_78>', '<extra_id_79>', '<extra_id_80>', '<extra_id_81>', '<extra_id_82>', '<extra_id_83>', '<extra_id_84>', '<extra_id_85>', '<extra_id_86>', '<extra_id_87>', '<extra_id_88>', '<extra_id_89>', '<extra_id_90>', '<extra_id_91>', '<extra_id_92>', '<extra_id_93>', '<extra_id_94>', '<extra_id_95>', '<extra_id_96>', '<extra_id_97>', '<extra_id_98>', '<extra_id_99>', '<extra_id_100>', '<extra_id_101>', '<extra_id_102>', '<extra_id_103>', '<extra_id_104>', '<extra_id_105>', '<extra_id_106>', '<extra_id_107>', '<extra_id_108>', '<extra_id_109>', '<extra_id_110>', '<extra_id_111>', '<extra_id_112>', '<extra_id_113>', '<extra_id_114>', '<extra_id_115>', '<extra_id_116>', '<extra_id_117>', '<extra_id_118>', '<extra_id_119>', '<extra_id_120>', '<extra_id_121>', '<extra_id_122>', '<extra_id_123>', '<extra_id_124>']}, clean_up_tokenization_spaces=True),  added_tokens_decoder={\n",
       "\t0: AddedToken(\"<pad>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t1: AddedToken(\"</s>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t2: AddedToken(\"<unk>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t256: AddedToken(\"<extra_id_0>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t257: AddedToken(\"<extra_id_1>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t258: AddedToken(\"<extra_id_2>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t259: AddedToken(\"<extra_id_3>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t260: AddedToken(\"<extra_id_4>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t261: AddedToken(\"<extra_id_5>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t262: AddedToken(\"<extra_id_6>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t263: AddedToken(\"<extra_id_7>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t264: AddedToken(\"<extra_id_8>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t265: AddedToken(\"<extra_id_9>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t266: AddedToken(\"<extra_id_10>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t267: AddedToken(\"<extra_id_11>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t268: AddedToken(\"<extra_id_12>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t269: AddedToken(\"<extra_id_13>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t270: AddedToken(\"<extra_id_14>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t271: AddedToken(\"<extra_id_15>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t272: AddedToken(\"<extra_id_16>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t273: AddedToken(\"<extra_id_17>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t274: AddedToken(\"<extra_id_18>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t275: AddedToken(\"<extra_id_19>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t276: AddedToken(\"<extra_id_20>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t277: AddedToken(\"<extra_id_21>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t278: AddedToken(\"<extra_id_22>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t279: AddedToken(\"<extra_id_23>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t280: AddedToken(\"<extra_id_24>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t281: AddedToken(\"<extra_id_25>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t282: AddedToken(\"<extra_id_26>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t283: AddedToken(\"<extra_id_27>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t284: AddedToken(\"<extra_id_28>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t285: AddedToken(\"<extra_id_29>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t286: AddedToken(\"<extra_id_30>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t287: AddedToken(\"<extra_id_31>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t288: AddedToken(\"<extra_id_32>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t289: AddedToken(\"<extra_id_33>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t290: AddedToken(\"<extra_id_34>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t291: AddedToken(\"<extra_id_35>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t292: AddedToken(\"<extra_id_36>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t293: AddedToken(\"<extra_id_37>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t294: AddedToken(\"<extra_id_38>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t295: AddedToken(\"<extra_id_39>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t296: AddedToken(\"<extra_id_40>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t297: AddedToken(\"<extra_id_41>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t298: AddedToken(\"<extra_id_42>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t299: AddedToken(\"<extra_id_43>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t300: AddedToken(\"<extra_id_44>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t301: AddedToken(\"<extra_id_45>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t302: AddedToken(\"<extra_id_46>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t303: AddedToken(\"<extra_id_47>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t304: AddedToken(\"<extra_id_48>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t305: AddedToken(\"<extra_id_49>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t306: AddedToken(\"<extra_id_50>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t307: AddedToken(\"<extra_id_51>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t308: AddedToken(\"<extra_id_52>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t309: AddedToken(\"<extra_id_53>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t310: AddedToken(\"<extra_id_54>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t311: AddedToken(\"<extra_id_55>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t312: AddedToken(\"<extra_id_56>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t313: AddedToken(\"<extra_id_57>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t314: AddedToken(\"<extra_id_58>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t315: AddedToken(\"<extra_id_59>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t316: AddedToken(\"<extra_id_60>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t317: AddedToken(\"<extra_id_61>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t318: AddedToken(\"<extra_id_62>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t319: AddedToken(\"<extra_id_63>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t320: AddedToken(\"<extra_id_64>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t321: AddedToken(\"<extra_id_65>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t322: AddedToken(\"<extra_id_66>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t323: AddedToken(\"<extra_id_67>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t324: AddedToken(\"<extra_id_68>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t325: AddedToken(\"<extra_id_69>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t326: AddedToken(\"<extra_id_70>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t327: AddedToken(\"<extra_id_71>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t328: AddedToken(\"<extra_id_72>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t329: AddedToken(\"<extra_id_73>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t330: AddedToken(\"<extra_id_74>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t331: AddedToken(\"<extra_id_75>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t332: AddedToken(\"<extra_id_76>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t333: AddedToken(\"<extra_id_77>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t334: AddedToken(\"<extra_id_78>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t335: AddedToken(\"<extra_id_79>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t336: AddedToken(\"<extra_id_80>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t337: AddedToken(\"<extra_id_81>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t338: AddedToken(\"<extra_id_82>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t339: AddedToken(\"<extra_id_83>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t340: AddedToken(\"<extra_id_84>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t341: AddedToken(\"<extra_id_85>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t342: AddedToken(\"<extra_id_86>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t343: AddedToken(\"<extra_id_87>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t344: AddedToken(\"<extra_id_88>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t345: AddedToken(\"<extra_id_89>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t346: AddedToken(\"<extra_id_90>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t347: AddedToken(\"<extra_id_91>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t348: AddedToken(\"<extra_id_92>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t349: AddedToken(\"<extra_id_93>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t350: AddedToken(\"<extra_id_94>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t351: AddedToken(\"<extra_id_95>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t352: AddedToken(\"<extra_id_96>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t353: AddedToken(\"<extra_id_97>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t354: AddedToken(\"<extra_id_98>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t355: AddedToken(\"<extra_id_99>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t356: AddedToken(\"<extra_id_100>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t357: AddedToken(\"<extra_id_101>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t358: AddedToken(\"<extra_id_102>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t359: AddedToken(\"<extra_id_103>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t360: AddedToken(\"<extra_id_104>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t361: AddedToken(\"<extra_id_105>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t362: AddedToken(\"<extra_id_106>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t363: AddedToken(\"<extra_id_107>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t364: AddedToken(\"<extra_id_108>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t365: AddedToken(\"<extra_id_109>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t366: AddedToken(\"<extra_id_110>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t367: AddedToken(\"<extra_id_111>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t368: AddedToken(\"<extra_id_112>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t369: AddedToken(\"<extra_id_113>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t370: AddedToken(\"<extra_id_114>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t371: AddedToken(\"<extra_id_115>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t372: AddedToken(\"<extra_id_116>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t373: AddedToken(\"<extra_id_117>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t374: AddedToken(\"<extra_id_118>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t375: AddedToken(\"<extra_id_119>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t376: AddedToken(\"<extra_id_120>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t377: AddedToken(\"<extra_id_121>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t378: AddedToken(\"<extra_id_122>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t379: AddedToken(\"<extra_id_123>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "\t380: AddedToken(\"<extra_id_124>\", rstrip=True, lstrip=True, single_word=False, normalized=False, special=True),\n",
       "}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, ByT5Tokenizer\n",
    "\n",
    "tokenizer = ByT5Tokenizer()\n",
    "tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "SOS_token = 0\n",
    "EOS_token = 2\n",
    "\n",
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, dropout_p=0.1):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.embedding = nn.Embedding(input_size, hidden_size)\n",
    "        self.gru = nn.LSTM(hidden_size, hidden_size, num_layers = 2, batch_first=True)\n",
    "        self.dropout = nn.Dropout(dropout_p)\n",
    "\n",
    "    def forward(self, input):\n",
    "        embedded = self.dropout(self.embedding(input))\n",
    "        output, hidden = self.gru(embedded)\n",
    "        return output, hidden\n",
    "\n",
    "class DecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size):\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        self.embedding = nn.Embedding(output_size, hidden_size)\n",
    "        self.gru = nn.LSTM(hidden_size, hidden_size, num_layers = 2, batch_first=True)\n",
    "        self.out = nn.Linear(hidden_size, output_size)\n",
    "        \n",
    "    @property\n",
    "    def device(self):\n",
    "        return next(self.parameters()).device\n",
    "\n",
    "    def forward(self, encoder_outputs, encoder_hidden, target_tensor=None):\n",
    "        batch_size = encoder_outputs.size(0)\n",
    "        decoder_input = torch.empty(batch_size, 1, dtype=torch.long, device=self.device).fill_(SOS_token)\n",
    "        decoder_hidden = encoder_hidden\n",
    "        decoder_outputs = []\n",
    "\n",
    "        if target_tensor is not None:\n",
    "            maxlen = target_tensor.shape[1]\n",
    "        else:\n",
    "            maxlen = encoder_outputs.shape[1] * 2\n",
    "\n",
    "        for i in range(maxlen):\n",
    "            decoder_output, decoder_hidden  = self.forward_step(decoder_input, decoder_hidden)\n",
    "            decoder_outputs.append(decoder_output)\n",
    "\n",
    "            if target_tensor is not None:\n",
    "                # Teacher forcing: Feed the target as the next input\n",
    "                decoder_input = target_tensor[:, i].unsqueeze(1) # Teacher forcing\n",
    "            else:\n",
    "                # Without teacher forcing: use its own predictions as the next input\n",
    "                _, topi = decoder_output.topk(1)\n",
    "                decoder_input = topi.squeeze(-1).detach()  # detach from history as input\n",
    "\n",
    "        decoder_outputs = torch.cat(decoder_outputs, dim=1)\n",
    "        decoder_outputs = F.log_softmax(decoder_outputs, dim=-1)\n",
    "        return decoder_outputs, decoder_hidden, None # We return `None` for consistency in the training loop\n",
    "\n",
    "    def forward_step(self, input, hidden):\n",
    "        output = self.embedding(input)\n",
    "        output = F.relu(output)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "        output = self.out(output)\n",
    "        return output, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, dropout_p=0.1):\n",
    "        super(Model, self).__init__()\n",
    "        self.encoder = EncoderRNN(input_size, hidden_size, dropout_p=dropout_p)\n",
    "        self.decoder = DecoderRNN(hidden_size, input_size)\n",
    "    \n",
    "    @property\n",
    "    def device(self):\n",
    "        return next(self.parameters()).device\n",
    "    \n",
    "    def forward(self, input_tensor, target_tensor=None):\n",
    "        input_tensor = input_tensor.to(self.device)\n",
    "        if target_tensor is not None:\n",
    "            target_tensor = target_tensor.to(self.device)\n",
    "        encoder_out = self.encoder(input_tensor)\n",
    "        return self.decoder(*encoder_out, target_tensor)\n",
    "    \n",
    "    def greedy_decoder(self, input_tensor):\n",
    "        input_tensor = input_tensor.to(self.device)\n",
    "        with torch.no_grad():\n",
    "            encoder_out = self.encoder(input_tensor)\n",
    "            decoder_outputs, decoder_hidden, _ = self.decoder(*encoder_out)\n",
    "            return decoder_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(tokenizer.vocab_size, 1024)\n",
    "_ = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainable_parameters = [param for param in model.parameters() if param.requires_grad]\n",
    "trainer = torch.optim.Adam(trainable_parameters, lr = 5e-4)\n",
    "criterion = nn.NLLLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from sklearn.utils import shuffle\n",
    "\n",
    "with open('train_noisy_stem.json') as fopen:\n",
    "    noisy = json.load(fopen)\n",
    "    \n",
    "with open('train_stem.json') as fopen:\n",
    "    data = json.load(fopen)\n",
    "    \n",
    "X = data['X'] + noisy['X']\n",
    "Y = data['Y'] + noisy['Y']\n",
    "X, Y = shuffle(X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('test_noisy_stem.json') as fopen:\n",
    "    noisy = json.load(fopen)\n",
    "    \n",
    "with open('test_stem.json') as fopen:\n",
    "    data = json.load(fopen)\n",
    "    \n",
    "test_X = data['X'] + noisy['X']\n",
    "test_Y = data['Y'] + noisy['Y']\n",
    "test_X, test_Y = shuffle(test_X, test_Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "strings = [f'{s}' for s in X[:10]]\n",
    "targets = [f'{s}' for s in Y[:10]]\n",
    "input_ids = tokenizer(strings, padding = True, return_tensors = 'pt')['input_ids']\n",
    "targets = tokenizer(targets, padding = True, return_tensors = 'pt')['input_ids']\n",
    "labels = torch.clone(targets)\n",
    "labels[targets==0] = -100\n",
    "labels = labels.to(model.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = model(input_ids, targets)\n",
    "decoder_outputs = out[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(5.5493, device='cuda:0', grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = criterion(\n",
    "        decoder_outputs.view(-1, decoder_outputs.size(-1)),\n",
    "        labels.view(-1)\n",
    "    )\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.greedy_decoder(input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 4.208215250400826e-05, accuracy 1.0: 100%|█| 91780/91780 [32:01<00:00, 47.7\n",
      "loss 0.04311397671699524, accuracy 0.9874476987447699: 100%|█| 313/313 [00:01<00\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, accuracy: 0.9938737785713346, dev_predicted: 0.9892604581383452\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.0001810267276596278, accuracy 1.0: 100%|█| 91780/91780 [31:35<00:00, 48.4\n",
      "loss 0.029947953298687935, accuracy 0.9832635983263598: 100%|█| 313/313 [00:01<0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, accuracy: 0.9939510437542196, dev_predicted: 0.989727926721791\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.00016259029507637024, accuracy 1.0: 100%|█| 91780/91780 [31:46<00:00, 48.\n",
      "loss 0.015072748064994812, accuracy 0.99581589958159: 100%|█| 313/313 [00:01<00:\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 2, accuracy: 0.9939598115398169, dev_predicted: 0.9897694818555247\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.0002972305519506335, accuracy 1.0: 100%|█| 91780/91780 [32:44<00:00, 46.7\n",
      "loss 0.022413861006498337, accuracy 0.9916317991631799: 100%|█| 313/313 [00:01<0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 3, accuracy: 0.9940350594681625, dev_predicted: 0.9899317396449283\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.0013899446930736303, accuracy 1.0: 100%|█| 91780/91780 [31:28<00:00, 48.6\n",
      "loss 0.025860510766506195, accuracy 0.9874476987447699: 100%|█| 313/313 [00:01<0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 4, accuracy: 0.9938956442534749, dev_predicted: 0.988657477579811\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.038468532264232635, accuracy 0.9890909090909091:  23%|▏| 21373/91780 [07:\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_917743/3791248802.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     24\u001b[0m         \u001b[0mlabels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_ids\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     27\u001b[0m         \u001b[0mdecoder_outputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1499\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1500\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1502\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1503\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_917743/825417015.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_tensor, target_tensor)\u001b[0m\n\u001b[1;32m     13\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mtarget_tensor\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m             \u001b[0mtarget_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtarget_tensor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m         \u001b[0mencoder_out\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mencoder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_tensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     16\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecoder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mencoder_out\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_tensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1499\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1500\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1502\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1503\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_917743/1261260463.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m         \u001b[0membedded\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdropout\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0membedding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m         \u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhidden\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgru\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0membedded\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     16\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhidden\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1499\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1500\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1502\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1503\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/rnn.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input, hx)\u001b[0m\n\u001b[1;32m    810\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcheck_forward_args\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_sizes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    811\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mbatch_sizes\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 812\u001b[0;31m             result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,\n\u001b[0m\u001b[1;32m    813\u001b[0m                               self.dropout, self.training, self.bidirectional, self.batch_first)\n\u001b[1;32m    814\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "batch_size = 32\n",
    "epoch = 100\n",
    "\n",
    "best_dev_acc = -np.inf\n",
    "patient = 2\n",
    "current_patient = 0\n",
    "\n",
    "for e in range(epoch):\n",
    "    pbar = tqdm(range(0, len(X), batch_size))\n",
    "    losses = []\n",
    "    for i in pbar:\n",
    "        trainer.zero_grad()\n",
    "        x = X[i: i + batch_size]\n",
    "        y = Y[i: i + batch_size]\n",
    "        strings = [s for s in x]\n",
    "        targets = [s for s in y]\n",
    "        input_ids = tokenizer(strings, padding = True, return_tensors = 'pt')['input_ids']\n",
    "        targets = tokenizer(targets, padding = True, return_tensors = 'pt')['input_ids']\n",
    "        labels = torch.clone(targets)\n",
    "        labels[targets==0] = -100\n",
    "        labels = labels.to(model.device)\n",
    "        \n",
    "        out = model(input_ids, targets)\n",
    "        decoder_outputs = out[0]\n",
    "        \n",
    "        loss = criterion(\n",
    "                decoder_outputs.view(-1, decoder_outputs.size(-1)),\n",
    "                labels.view(-1)\n",
    "            )\n",
    "        loss.backward()\n",
    "        trainer.step()\n",
    "        \n",
    "        p = decoder_outputs.detach().cpu().numpy().argmax(axis = -1)\n",
    "        l = labels.detach().cpu().numpy()\n",
    "        mask = l != -100\n",
    "        accuracy = ((p == l)[mask]).mean()\n",
    "        losses.append(accuracy)\n",
    "        pbar.set_description(f\"loss {loss}, accuracy {accuracy}\")\n",
    "    \n",
    "    pbar = tqdm(range(0, len(test_X[:10000]), batch_size))\n",
    "    dev_predicted = []\n",
    "    for i in pbar:\n",
    "        x = test_X[i: i + batch_size]\n",
    "        y = test_Y[i: i + batch_size]\n",
    "        strings = [s for s in x]\n",
    "        targets = [s for s in y]\n",
    "        input_ids = tokenizer(strings, padding = True, return_tensors = 'pt')['input_ids']\n",
    "        targets = tokenizer(targets, padding = True, return_tensors = 'pt')['input_ids']\n",
    "        labels = torch.clone(targets)\n",
    "        labels[targets==0] = -100\n",
    "        labels = labels.to(model.device)\n",
    "        \n",
    "        out = model(input_ids, targets)\n",
    "        decoder_outputs = out[0]\n",
    "        \n",
    "        loss = criterion(\n",
    "                decoder_outputs.view(-1, decoder_outputs.size(-1)),\n",
    "                labels.view(-1)\n",
    "            )\n",
    "        \n",
    "        p = decoder_outputs.detach().cpu().numpy().argmax(axis = -1)\n",
    "        l = labels.detach().cpu().numpy()\n",
    "        mask = l != -100\n",
    "        accuracy = ((p == l)[mask]).mean()\n",
    "        dev_predicted.append(accuracy)\n",
    "        pbar.set_description(f\"loss {loss}, accuracy {accuracy}\")\n",
    "    \n",
    "    dev_predicted = np.mean(dev_predicted)\n",
    "    \n",
    "    print(f'epoch: {e}, accuracy: {np.mean(losses)}, dev_predicted: {dev_predicted}')\n",
    "    \n",
    "    if dev_predicted >= best_dev_acc:\n",
    "        best_dev_acc = dev_predicted\n",
    "        current_patient = 0\n",
    "    else:\n",
    "        current_patient += 1\n",
    "    \n",
    "    if current_patient >= patient:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# strings = ['kelemahan', 'gaykan']\n",
    "# input_ids = tokenizer(strings, padding = True, return_tensors = 'pt')['input_ids']\n",
    "# out = model.greedy_decoder(input_ids)\n",
    "# tokenizer.batch_decode(out.detach().cpu().numpy().argmax(axis = -1), skip_special_tokens = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_cer(actual, hyp):\n",
    "    \"\"\"\n",
    "    Calculate CER using `python-Levenshtein`.\n",
    "    \"\"\"\n",
    "    import Levenshtein as Lev\n",
    "\n",
    "    actual = actual.replace(' ', '')\n",
    "    hyp = hyp.replace(' ', '')\n",
    "    return Lev.distance(actual, hyp) / len(actual)\n",
    "\n",
    "\n",
    "def calculate_wer(actual, hyp):\n",
    "    \"\"\"\n",
    "    Calculate WER using `python-Levenshtein`.\n",
    "    \"\"\"\n",
    "    import Levenshtein as Lev\n",
    "\n",
    "    b = set(actual.split() + hyp.split())\n",
    "    word2char = dict(zip(b, range(len(b))))\n",
    "\n",
    "    w1 = [chr(word2char[w]) for w in actual.split()]\n",
    "    w2 = [chr(word2char[w]) for w in hyp.split()]\n",
    "\n",
    "    return Lev.distance(''.join(w1), ''.join(w2)) / len(actual.split())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████| 29551/29551 [05:47<00:00, 84.98it/s]\n"
     ]
    }
   ],
   "source": [
    "cers, wers = [], []\n",
    "\n",
    "batch_size = 16\n",
    "for i in tqdm(range(0, len(test_X), batch_size)):\n",
    "    x = test_X[i: i + batch_size]\n",
    "    y = test_Y[i: i + batch_size]\n",
    "    strings = [s for s in x]\n",
    "    input_ids = tokenizer(strings, padding = True, return_tensors = 'pt')['input_ids']\n",
    "    out = model.greedy_decoder(input_ids)\n",
    "    argmax = out.detach().cpu().numpy().argmax(axis = -1)\n",
    "    results = []\n",
    "    for a in argmax:\n",
    "        results_ = []\n",
    "        for a_ in a:\n",
    "            if a_ == tokenizer.eos_token_id:\n",
    "                break\n",
    "            else:\n",
    "                results_.append(a_)\n",
    "        results.append(results_)\n",
    "    out = tokenizer.batch_decode(results)\n",
    "    for no, y_ in enumerate(y):\n",
    "        cers.append(calculate_cer(y_, out[no]))\n",
    "        wers.append(calculate_wer(y_, out[no]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.03244944985606283, 0.06760166327670675)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.mean(cers), np.mean(wers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'lstm-2-1024.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/stem-lstm-1024/commit/578d2182b0f60d6d5fbd0c81357119fe04ac5782', commit_message='Upload tokenizer', commit_description='', oid='578d2182b0f60d6d5fbd0c81357119fe04ac5782', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.push_to_hub('mesolitica/stem-lstm-1024')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "from malaya_boilerplate.huggingface import upload_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "409 Client Error: Conflict for url: https://huggingface.co/api/repos/create (Request ID: Root=1-6527ac3f-7cd4856714583b6f1bf0fb85;fa35b37b-5104-4a47-9fd9-f12623d301f8)\n",
      "\n",
      "You already created this model repo\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76cea57d9b954b589ac8d19a5ce9fd92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "lstm-2-1024.pt:   0%|          | 0.00/137M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "upload_dict('stem-lstm-1024', {'lstm-2-1024.pt': 'model.pt'}, username = 'mesolitica')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
